# -*- coding: utf-8 -*-
"""
计算 mAp and PR
测试结果存在result/result.txt中，每一行的检测信息格式为：
image_name class_id confidence xmin ymin xmax ymax
"""
import os
import sys
import argparse
import xml.etree.ElementTree as ET
import numpy as np
import matplotlib.pyplot as plt


# import cPickle as pk


class Box:
    """
    bbox
    """
    def __init__(self, x, y, w, h):
        self.x = x
        self.y = y
        self.w = w
        self.h = h


def parse_ground_truth(root_dir, file_list, classes, filter_difficult):
    """
    parse ground truth from xml files.
    :param root_dir: voc root path
    :param file_list: test.txt
    :param classes: class labels
    :param filter_difficult: filter out difficult targets if True
    :return: ground_true_per_img: list of object informations per image;
             object_num: all the object numbers in test set
    """
    ground_true_per_img = {}
    object_num = 0
    test_files = open(file_list, 'r').readlines()
    for test_file in test_files:
        xml_file = os.path.join(root_dir, 'Annotations/{}.xml'.format(test_file.strip()))
        tree = ET.parse(xml_file)
        root = tree.getroot()
        object_info = []

        for obj in root.iter('object'):
            difficult = obj.find('difficult').text
            object_class = obj.find('name').text

            if object_class in classes and (filter_difficult and int(difficult) == 1):
                continue

            xml_box = obj.find('bndbox')
            rect = (float(xml_box.find('xmin').text),
                    float(xml_box.find('ymin').text),
                    float(xml_box.find('xmax').text),
                    float(xml_box.find('ymax').text))
            object_info.append({'object_class': object_class, 'rect': rect})
            object_num += 1
        ground_true_per_img[test_file.strip()] = object_info
    return ground_true_per_img, object_num


def parse_detect_result(detect_result_file):
    """

    :param detect_result_file:
    :return:
    """

    detect_result_lst = []

    detect_results = open(detect_result_file, 'r').readlines()
    for result in detect_results:
        detect_info = result.strip().split()

        image_name = detect_info[0]
        class_name = detect_info[1]
        confidence = float(detect_info[2])

        left = float(detect_info[3])
        top = float(detect_info[4])
        right = float(detect_info[5])
        bottom = float(detect_info[6])
        rect = (left, top, right, bottom)
        detect_result_lst.append(
            {'image_name': image_name,
             'class_name': class_name,
             'confidence': confidence,
             'rect': rect})

    return detect_result_lst

"""

def sort_comp_desc(a, b):
    if a['confidence'] > b['confidence']:
        return -1
    elif a['confidence'] < b['confidence']:
        return 1
    else:
        return 0
"""


def overlap(center_1, len_1, center_2, len_2):
    """
    compute axis overlap
    :param center_1: box1 center on one axis
    :param len_1:    box1 length on one axis
    :param center_2: box2 center on one axis
    :param len_2:    box2 length on one axis
    :return: overlap
    """
    len1_half = len_1 / 2
    len2_half = len_2 / 2

    axis_min = max(center_1 - len1_half, center_2 - len2_half)
    axis_max = min(center_1 + len1_half, center_2 + len2_half)

    return axis_max - axis_min


def box_intersection(a, b):
    """
    compute intersection
    :param a: Box a
    :param b: Box b
    :return: intersection
    """
    overlap_w = overlap(a.x, a.w, b.x, b.w)
    overlap_h = overlap(a.y, a.h, b.y, b.h)
    if overlap_w <= 0 or overlap_h <= 0:
        return 0

    area = overlap_w * overlap_w
    return area


def box_union(a, b):
    """
    compute union
    :param a: Box a
    :param b: Box b
    :return: union
    """
    intersection = box_intersection(a, b)
    union = a.w * a.h + b.w * b.h - intersection
    return union


def box_iou(a, b):
    """
    compute iou
    :param a: Box a
    :param b: Box b
    :return: iou
    """
    return box_intersection(a, b) / box_union(a, b)


def cal_ap(recall, precision):
    """
    cal ap
    :param recall:
    :param precision:
    :return:
    """
    ap_val = 0.0
    eps = 1e-6
    assert len(recall) == len(precision)
    length = len(recall)
    cur_precision = precision[length - 1]
    cur_recall = recall[length - 1]

    for i in range(0, length - 1)[::-1]:
        cur_precision = max(precision[i], cur_precision)
        if abs(recall[i] - cur_recall) > eps:
            ap_val += cur_precision * abs(recall[i] - cur_recall)

        cur_recall = recall[i]

    return ap_val


def cal_ap_11point(recall, precision):
    ap_val = 0.0
    eps = 1e-6
    assert len(recall) == len(precision)
    num = len(recall)
    max_precs = np.zeros(num + 1)

    start_idx = num - 1
    for j in range(0, 11)[::-1]:
        for i in range(0, start_idx + 1)[::-1]:
            if recall[i] < (j / 10.0):
                start_idx = i
                if j > 0:
                    max_precs[j - 1] = max_precs[j]
                break
            else:
                if max_precs[j] < precision[i]:
                    max_precs[j] = precision[i]

    for j in range(0, 11):
        ap_val += max_precs[j] / 11.0

    return ap_val


def count_ground_truth_object(ground_truth_hashmap, classes):
    """
    count object numbers belong to 'classes'
    :param ground_truth_hashmap: ground truth
    :param classes: class names the object belongs to.
    :return: total object numbers
    """
    counter = 0
    for key_v in ground_truth_hashmap.keys():
        for obj_ground_truth in ground_truth_hashmap[key_v]:
            if obj_ground_truth['object_class'] in classes:
                counter += 1
    return counter


def cal_precision_ap(ground_truth_hashmap, detect_result, classes, iou_thr):
    """

    :param ground_truth_hashmap:
    :param detect_result:
    :param classes:
    :param iou_thr:
    :return:
    """
    sorted_detect_result = sorted(detect_result,
                                  key=lambda detect_result: detect_result['confidence'],
                                  reverse=True)
    tp = np.zeros(len(sorted_detect_result))
    fp = np.zeros(len(sorted_detect_result))

    # count object numbers belong to [classes]
    positive_num = count_ground_truth_object(ground_truth_hashmap, classes)

    for result_idx, object_info in enumerate(detect_result):
        image_name = object_info['image_name']
        image_name = os.path.split(image_name)[-1].split(".")[0]

        if image_name not in ground_truth_hashmap:
            print("image {}.jpg is not in test list".format(image_name))
            continue
        ground_truth = ground_truth_hashmap[image_name]
        iou_max = 0
        max_idx = -1

        detect_x = object_info['rect'][0]
        detect_y = object_info['rect'][1]
        detect_w = object_info['rect'][2] - detect_x
        detect_h = object_info['rect'][3] - detect_y
        box_detect = Box(detect_x, detect_y, detect_w, detect_h)

        for ground_truth_object_idx, object_ground_truth in enumerate(ground_truth):
            if not object_ground_truth['object_class'] in classes or 'used' in object_ground_truth:
                continue

            ground_truth_x = object_ground_truth['rect'][0]
            ground_truth_y = object_ground_truth['rect'][1]
            ground_truth_w = object_ground_truth['rect'][2] - ground_truth_x
            ground_truth_h = object_ground_truth['rect'][3] - ground_truth_y
            box_ground_truth = Box(ground_truth_x, ground_truth_y, ground_truth_w, ground_truth_h)

            iou = box_iou(box_detect, box_ground_truth)
            if iou > iou_max:
                iou_max = iou
                max_idx = ground_truth_object_idx

        if iou_max > iou_thr:
            if 'used' in ground_truth[max_idx]:
                fp[result_idx] = 1
            else:
                tp[result_idx] = 1
                ground_truth[max_idx]['used'] = 1
        else:
            fp[result_idx] = 1

    print('del used key!')
    del_nr = 0
    for ground_truth_keys in ground_truth_hashmap.keys():
        for obj_idx, image_ground_truth_info in \
                enumerate(ground_truth_hashmap[ground_truth_keys]):
            # if(a_gt.has_key('used')):
            if 'used' in image_ground_truth_info:
                ground_truth_hashmap[ground_truth_keys][obj_idx].pop('used')
                del_nr += 1

    print('del key nr=%d' % del_nr)

    tp = tp.cumsum()
    fp = fp.cumsum()

    recall = tp / positive_num  # type: float
    precision = tp / (tp + fp)

    ap = cal_ap(recall, precision)

    return recall, precision, ap


if __name__ == '__main__':
    try:
        ROOT_DIR = sys.argv[1]
        IOU_THRES = float(sys.argv[2])
    except:
        PARSER = argparse.ArgumentParser(description="arguments")
        PARSER.add_argument("--root_dir", type=str,
                            default="/home/chujie/PycharmProjects/data_tool/data/voc2012",
                            help="VOC dataset root directory")
        PARSER.add_argument("--iou_thres", type=float, default=0.2,
                            help="iou threshold")
        ARGS = PARSER.parse_args()
        ROOT_DIR = ARGS.root_dir
        IOU_THRES = ARGS.iou_thres
    LABELS = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus",
              "car", "cat", "chair", "cow", "diningtable", "dog", "horse",
              "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
    # DETECT_RESULT_FILE = os.path.join(ROOT_DIR, "result/result.txt")
    TEST_FILE_LIST = os.path.join(ROOT_DIR, "ImageSets/Main/test.txt")

    # ground_truth_hash_map: [{"image_name":
    # {"obj_class": class, "rect": (xmin, ymin, xmax, ymax)} }]
    ground_truth_hash_map, total_object_num = \
        parse_ground_truth(ROOT_DIR, TEST_FILE_LIST, LABELS[14], True)
    # detect_result_lst: [{'image_id':
    # img_id, 'class_id': class_id, 'confidence': confidence, 'rect': rect}]
    # rect: [xmin, ymin, xmax, ymax]
    DETECT_RESULT_LIST = parse_detect_result(os.path.join(ROOT_DIR, "result/result.txt"))
    RECALL, PRECISION, AP = cal_precision_ap(ground_truth_hash_map,
                                             DETECT_RESULT_LIST, LABELS[14], IOU_THRES)

    # show
    per_line, = plt.plot(RECALL, PRECISION, 'b-')
    per_line.set_label('yolo_%s:AP=%.3f' % (LABELS[14], AP))

    # per_line2, = plt.plot(recall2,precision2,'r-')
    # per_line2.set_label('ssd_%s:AP=%.3f'%(classes[0],ap2))

    plt.plot([i / 1000.0 for i in range(1, 1001)], [i / 1000.0 for i in range(1, 1001)], 'y--')

    plt.axis([0, 1.2, 0, 1.2])
    plt.xlabel('recall')
    plt.ylabel('precision')
    plt.grid()

    plt.legend()
    plt.title('PR')
    plt.show()
